#include <stdio.h>
#include <pcap.h>
#include <time.h>
#include <stdint.h>
#include <pthread.h>
#include <stdlib.h>
#include "/home/haiyang/wendang/Cprogrammer/include/packet_processing.h"
#include "/home/haiyang/wendang/Cprogrammer/include/stats_saving.h"

#define SNAP_LEN 1518
#define BUFFER_SIZE 1024

void *packet_capture_thread(void *arg);
void packet_handler(uint8_t *user_data, const struct pcap_pkthdr *pkthdr, const uint8_t *packet);

int main() {
    pthread_t thread_id;
    TrafficStats stats;
    init_traffic_stats(&stats);

    // 创建线程进行抓包和统计
    if (pthread_create(&thread_id, NULL, packet_capture_thread, &stats) != 0) {
        perror("Failed to create thread");
        return 1;
    }

    // 等待线程结束traffic_monitor
    if (pthread_join(thread_id, NULL) != 0) {
        perror("Failed to join thread");
        return 1;
    }

    // 保存统计报告
    save_stats_to_file(&stats);

    return 0;
}

void *packet_capture_thread(void *arg) {
    TrafficStats *stats = (TrafficStats *)arg;
    char errbuf[PCAP_ERRBUF_SIZE];
    pcap_t *handle;
    struct bpf_program fp;
    char filter_exp[] = "ip";
    bpf_u_int32 net;

    pcap_if_t *alldevs;
    pcap_if_t *d;
    int i = 0;

    // 获取所有可用的网络接口
    if (pcap_findalldevs(&alldevs, errbuf) == -1) {
        fprintf(stderr, "Error in pcap_findalldevs: %s\n", errbuf);
        pthread_exit((void *)2);
    }

    // 打印所有可用的网络接口
    for (d = alldevs; d != NULL; d = d->next) {
        printf("%d. %s", ++i, d->name);
        if (d->description)
            printf(" (%s)\n", d->description);
        else
            printf(" (No description available)\n");
    }

    if (i == 0) {
        printf("\nNo interfaces found! Make sure libpcap is installed.\n");
        pcap_freealldevs(alldevs);
        pthread_exit((void *)2);
    }

    // 选择第一个可用的网络接口
    handle = pcap_open_live(alldevs->name, SNAP_LEN, 1, 1000, errbuf);
    if (handle == NULL) {
        fprintf(stderr, "Couldn't open device %s: %s\n", alldevs->name, errbuf);
        pcap_freealldevs(alldevs);
        pthread_exit((void *)2);
    }

    // 释放设备列表
    pcap_freealldevs(alldevs);

    // 编译过滤规则
    if (pcap_compile(handle, &fp, filter_exp, 0, net) == -1) {
        fprintf(stderr, "Couldn't parse filter %s: %s\n", filter_exp, pcap_geterr(handle));
        pcap_close(handle);
        pthread_exit((void *)2);
    }

    // 应用过滤规则
    if (pcap_setfilter(handle, &fp) == -1) {
        fprintf(stderr, "Couldn't install filter %s: %s\n", filter_exp, pcap_geterr(handle));
        pcap_freecode(&fp);
        pcap_close(handle);
        pthread_exit((void *)2);
    }



    //修改时间
    //单位秒s
    time_t start_time = time(NULL);
    time_t end_time = start_time + 60;
    printf("Start time: %ld\n", start_time);
    printf("End time: %ld\n", end_time);

    while (time(NULL) < end_time) {
        time_t current_time = time(NULL);
        printf("Current time: %ld\n", current_time);
        pcap_dispatch(handle, -1, packet_handler, (uint8_t *)stats);
    }
    pcap_freecode(&fp);
    pcap_close(handle);
    pthread_exit(NULL);
}

void packet_handler(uint8_t *user_data, const struct pcap_pkthdr *pkthdr, const uint8_t *packet) {
    TrafficStats *stats = (TrafficStats *)user_data;
    TreeNode *root = build_tree();
    process_packet_through_tree(root, packet, pkthdr->len, stats);
    free_tree(root);
}